// Copyright (c) 2024 Huawei Technologies Co., Ltd.
// openFuyao is licensed under Mulan PSL v2.
// You can use this software according to the terms and conditions of the Mulan PSL v2.
// You may obtain a copy of Mulan PSL v2 at:
//          http://license.coscl.org.cn/MulanPSL2
// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
// EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
// MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
// See the Mulan PSL v2 for more details.

package lm

import (
	"fmt"
	"strings"

	"github.com/urfave/cli/v2"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	kubernetes2 "k8s.io/client-go/kubernetes"
	"k8s.io/klog/v2"

	"openfuyao.com/npu-feature-discovery/internal/kubernetes"
	"openfuyao.com/npu-feature-discovery/internal/utils"
)

func newAscendLabeler(c *cli.Context) (Labeler, error) {
	_, node, err := getNode(c)
	if err != nil {
		return nil, err
	}

	runtimeLabeler, err := newRuntimeLabeler(node)
	if err != nil {
		return nil, err
	}

	var labelers []Labeler
	if runtimeLabeler != nil {
		labelers = append(labelers, runtimeLabeler)

		ascendPresenceLabeler, hasAscend := newAscendPresenceLabeler(node)

		if hasAscend {
			labelers = append(labelers, ascendPresenceLabeler)

			selectorLabeler := newSelectorLabeler(node)
			labelers = append(labelers, selectorLabeler)

			hasDevicePlugin := checkDevicePlugin(node)
			if hasDevicePlugin {
				hostArchLabeler := newHostArchLabeler(node)
				labelers = append(labelers, hostArchLabeler)

				acceleratorLabeler := newAcceleratorLabeler(node)
				labelers = append(labelers, acceleratorLabeler)

				acceleratorTypeLabeler, err := newAcceleratorTypeLabeler(node)
				if err != nil {
					klog.Error(err)
				} else {
					labelers = append(labelers, acceleratorTypeLabeler)
				}
			}
		}
	}

	return merge(labelers...), nil
}

func checkDevicePlugin(node *corev1.Node) bool {
	for resourceName, resourceQuantity := range node.Status.Capacity {
		if strings.HasPrefix(resourceName.String(), "huawei.com/") {
			if resourceQuantity.Sign() > 0 {
				return true
			}
		}
	}
	klog.Info("No device plugin resources found or resource quantity is 0.")
	return false
}

func getNode(c *cli.Context) (kubernetes2.Interface, *corev1.Node, error) {
	client, err := kubernetes.GetKubernetesClient()
	if err != nil {
		return nil, nil, fmt.Errorf("failed to get Kubernetes client: %v", err)
	}

	nodeName, err := utils.GetEnvVar("NODE_NAME")
	if err != nil {
		return nil, nil, err
	}

	node, err := client.CoreV1().Nodes().Get(c.Context, nodeName, metav1.GetOptions{})
	if err != nil {
		return nil, nil, fmt.Errorf("failed to get node %s: %w", nodeName, err)
	}

	return client, node, nil
}
